/**************************************************************************
* Copyright(c) 1998-1999, ALICE Experiment at CERN, All rights reserved. *
*                                                                        *
* Author: The ALICE Off-line Project.                                    *
* Contributors are mentioned in the code where appropriate.              *
*                                                                        *
* Permission to use, copy, modify and distribute this software and its   *
* documentation strictly for non-commercial purposes is hereby granted   *
* without fee, provided that the above copyright notice appears in all   *
* copies and that both the copyright notice and this permission notice   *
* appear in the supporting documentation. The authors make no claims     *
* about the suitability of this software for any purpose. It is          *
* provided "as is" without express or implied warranty.                  *
**************************************************************************/
//
//Xian-Guo Lu <lu@physi.uni-heidelberg.de>
//

#ifndef ALIANALYSISTASKSE_H
#include "AliAnalysisTaskSE.h"
#endif

#include <TTreeStream.h>
#include "TH1I.h"

#include "AliAnalysisManager.h"
#include "AliAODHandler.h"
#include "AliAODInputHandler.h"
#include "AliCentrality.h"
#include "AliTPCdEdxInfo.h"

typedef AliAODTrack typetrack;

class AliAnalysisTaskAODpkRaa : public AliAnalysisTaskSE
{
 public:
  
  AliAnalysisTaskAODpkRaa();
  ~AliAnalysisTaskAODpkRaa();
  void UserExec(Option_t *);
  void UserCreateOutputObjects();

 private:
  Bool_t CutInput();
  Bool_t CutEvent();
  void CutTrack(Int_t &nsel, typetrack *outtrk[]);
  TString GetRunInfo();

  AliAODEvent *fEvent; 
  TTreeStream *fStream;
  TTree *fTree;

  TList *fList;
  TH1I * fCuts;
  TH1I *fNsel;
  TH1D *fCen;
  TH1D *fEta;
  TH1I *fFilterBit;

  ClassDef(AliAnalysisTaskAODpkRaa, 0);
};


AliAnalysisTaskAODpkRaa::AliAnalysisTaskAODpkRaa():
  AliAnalysisTaskSE("AODpkRaa")
  , fEvent(0x0)
  , fStream(0x0)
  , fTree(0x0)
  , fList(0x0)
  , fCuts(0x0)
  , fNsel(0x0)
  , fCen(0x0)
  , fEta(0x0)
  , fFilterBit(0x0)
{
  //
  //
  //

  DefineOutput(1, TTree::Class());
  DefineOutput(2, TList::Class());

  fStream = new TTreeStream("tree");
  fTree=fStream->GetTree();

  fList = new TList;
  fList->SetOwner(kTRUE);

  fCuts = new TH1I("cuts","", 30,0,30);
  fNsel = new TH1I("nsel","",100,0,100);
  fCen = new TH1D("cen","",110,-10,100);
  fEta = new TH1D("eta","",100,-1,1);
  fFilterBit = new TH1I("filterbit","",20,0,20);

  fList->Add(fCuts);
  fList->Add(fNsel);
  fList->Add(fCen);
  fList->Add(fEta);
  fList->Add(fFilterBit);
}

AliAnalysisTaskAODpkRaa::~AliAnalysisTaskAODpkRaa()
{
  //
  // Destructor
  //

  /* can not be deleted in chain, they are in input handler!
  delete fEvent;
  */

  //delete fStream;
  delete fList;
}

void AliAnalysisTaskAODpkRaa::UserCreateOutputObjects()
{
  //
  //
  //

  PostData(1, fTree);
  PostData(2, fList);
}

void AliAnalysisTaskAODpkRaa::UserExec(Option_t *)
{
  //
  //Execute
  // 

  fCuts->Fill(0); //================================

  if(!CutInput())
    return;

  fCuts->Fill(1); //================================

  if(!CutEvent())
    return;

  fCuts->Fill(2); //================================

  const Int_t ntrk0 = fEvent->GetNumberOfTracks();
  typetrack * trks[ntrk0];
  for(Int_t ii=0; ii<ntrk0; ii++){
    trks[ii] = 0x0;
  }
  Int_t nsel = 0;
  CutTrack(nsel, trks);
  fNsel->Fill(nsel);
  if(!nsel)
    return;

  fCuts->Fill(3); //================================

  Float_t v0cen = -999;
  const AliCentrality *evtCentrality = fEvent->GetCentrality();
  if(evtCentrality){
    v0cen = evtCentrality->GetCentralityPercentile("V0M");
  }

  //apply centrality cut here to accelerate
  fCen->Fill(v0cen);

  if(GetRunInfo().Contains("PbPb")){
    if(v0cen<0 || v0cen>5){
      return;
    }
  }

  fCuts->Fill(4); //================================

  //-------------

  for(Int_t itrk=0; itrk<nsel ; itrk++){
    fCuts->Fill(20);

    typetrack * trackptr = trks[itrk];

    Float_t mom = trackptr->P();
    Int_t charge = trackptr->Charge();
    Float_t pt = trackptr->Pt();

    Float_t eta = trackptr->Eta();
    fEta->Fill(eta);
    if(TMath::Abs(eta)<0.6 || TMath::Abs(eta)>0.8)
      continue;

    fCuts->Fill(21);

    Int_t tpcncross = -999;
    AliTPCdEdxInfo *tpcinfo = trackptr->GetTPCdEdxInfo();
    Double32_t infosignal[4];
    Char_t infoncl[3];
    Char_t infonrows[3];
    if(tpcinfo){
      tpcinfo->GetTPCSignalRegionInfo(infosignal, infoncl, infonrows);
      tpcncross = infonrows[0]+infonrows[1]+infonrows[2];
    }

    Float_t tpcsig  = trackptr->GetTPCsignal();
    Float_t tpcip = trackptr->GetTPCmomentum();
    Int_t tpcnf = trackptr->GetTPCNclsF();
    Int_t tpcpidn = trackptr->GetTPCsignalN();
    
    /*
      UShort_t   AliESDtrack::GetTPCnclsS(Int_t i0,Int_t i1) const{
      return  fTPCSharedMap.CountBits(i0)-fTPCSharedMap.CountBits(i1);
      }

      UShort_t   AliESDtrack::GetTPCncls(Int_t i0,Int_t i1) const{
      return  fTPCClusterMap.CountBits(i0)-fTPCClusterMap.CountBits(i1);
      }
    */

    const TBits tpcClusterMap = trackptr->GetTPCClusterMap();
    Int_t tpcncls = tpcClusterMap.CountBits(0)-tpcClusterMap.CountBits(159);

    const TBits tpcSharedMap = trackptr->GetTPCSharedMap();
    Int_t tpcns = tpcSharedMap.CountBits(0)-tpcSharedMap.CountBits(159);

    (*fStream)<<"tree"<<
      "v0cen="<<v0cen<<

      "mom="<<mom<<
      "charge="<<charge<<
      "pt="<<pt<<
      "eta="<<eta<<

      "tpcsig="<<tpcsig<<
      "tpcip="<<tpcip<<
      "tpcpidn="<<tpcpidn<<
      "tpcncls="<<tpcncls<<
      "tpcns="<<tpcns<<
      "tpcnf="<<tpcnf<<
      "tpcncross="<<tpcncross;

    if(tpcinfo){
      for(Int_t itmp=0; itmp<4; itmp++){
        (*fStream)<<"tree"<<
          Form("infosignal%d=",itmp)<<infosignal[itmp];
      }
      
      for(Int_t itmp=0; itmp<3; itmp++){
        (*fStream)<<"tree"<<
          Form("infoncl%d=",itmp)<<infoncl[itmp];
      }
      
      for(Int_t itmp=0; itmp<3; itmp++){
        (*fStream)<<"tree"<<
          Form("infonrows%d=",itmp)<<infonrows[itmp];
      }
    }

    (*fStream)<<"tree"<<
      "\n";
    
    fCuts->Fill(22);
  } 
  fCuts->Fill(5); //================================
}

TString AliAnalysisTaskAODpkRaa::GetRunInfo()
{
  const Int_t run = fEvent->GetRunNumber();
  //checked on E-log ->
  //10d:        -> 126460
  //10e: 126461 -> 130930
  //10f: 130931 ->

  //10g:        -> 136781
  //10h: 136782 -> 139846
  //11a: 139847 ->

  //11f:        -> 165771 
  //11h: 165772 -> 170718
  //12a: 170719 ->
  //<-

  if(run >= 126461 && run <= 130930){
    return "pp10e";
  }
  else if(run >= 136782 && run <= 139846){
    return "PbPb10h";
  }
  else if(run >= 165772 && run <= 170718){
    return "PbPb11h";
  }
  else{
    printf("AliAnalysisTaskAODpkRaa::GetRunInfo wrong run! %d\n", run);exit(1);
  }
}

Bool_t AliAnalysisTaskAODpkRaa::CutInput()
{
  //
  //general protection
  //
  fEvent = 0x0;

  AliVEventHandler * evtH = AliAnalysisManager::GetAnalysisManager()->GetInputEventHandler();
  if(!evtH){
    printf("AliAnalysisTaskAODpkRaa::CutInput error No evtH\n"); exit(1);
  }

  if(evtH->InheritsFrom("AliAODInputHandler")){
    fEvent =  ((AliAODInputHandler*)evtH)->GetEvent();
  }
  else if(evtH->InheritsFrom("AliAODHandler")){
    fEvent = ((AliAODHandler*)evtH)->GetAOD();
  }
  else{
    printf("AliAnalysisTaskAODpkRaa::CutInput error No aodH type %s\n", evtH->GetDataType()); exit(1);
  }

  if(!fEvent){
    printf("AliAnalysisTaskAODpkRaa::CutInput error No AODEvent  AODInputHandler %d AODHandler %d\n", evtH->InheritsFrom("AliAODInputHandler"), evtH->InheritsFrom("AliAODHandler")); exit(1);
  }

  //in unit of kg
  const Float_t dm = fEvent->GetMagneticField();
  if( TMath::Abs(TMath::Abs(dm)- 5) > 0.1 ){
    printf("AliAnalysisTaskAODpkRaa::CutInput error strange Bfield! %f\n", dm); exit(1);
  }

  return kTRUE;
}


Bool_t AliAnalysisTaskAODpkRaa::CutEvent()
{
  Int_t id = 10;
  fCuts->Fill(id++);
  
  const AliAODVertex *primVertex = fEvent->GetPrimaryVertex();
  if(!primVertex){
    return kFALSE;
  }

  fCuts->Fill(id++);
  
  //3 types of vertex: Tracks, SPD, TPC
  //do not use TPC vertex
  const TString pvname(primVertex->GetName());
  if(pvname=="TPCVertex"){
    return kFALSE;
  }

  fCuts->Fill(id++);
  
  if(primVertex->GetNContributors()<1){
    return kFALSE;
  }

  fCuts->Fill(id++);
  
  if(TMath::Abs( primVertex->GetZ() )>10){
    return kFALSE;
  }

  fCuts->Fill(id++);
  
  return kTRUE;
}

void AliAnalysisTaskAODpkRaa::CutTrack(Int_t &nsel,  typetrack *outtrk[])
{
  /*
  //___________________________________________________________________________
  //___________________________________________________________________________
  //10h_AODGSI /hera/alice/rromita/trunk_AOD/aod_AOD/AddTask_aod_ESDFilter_PbPb.C
  // standard cuts with tight DCA cut, using cluster cut instead of crossed rows (a la 2010 default)                                                          
  AliESDtrackCuts* esdTrackCutsH2Cluster = AliESDtrackCuts::GetStandardITSTPCTrackCuts2010(kTRUE, 0);
  // 1024, 1<< 10                                                              
  trackFilter->AddCuts(esdTrackCutsH2Cluster);
  //___________________________________________________________________________
  //___________________________________________________________________________
  //10h golden cut
  //https://alimonitor.cern.ch/users/download.jsp?view=true&path=/alice/cern.ch/user/a/alidaq/AOD/AOD086/AddTaskESDFilter.C
  // standard cuts also used in R_AA analysis
  //   "Global track RAA analysis QM2011 + Chi2ITS<36";
  // 1024 1<<10        
  AliESDtrackCuts* esdTrackCuts = AliESDtrackCuts::GetStandardITSTPCTrackCuts2010(kTRUE,1);
  esdTrackCuts->SetMinNCrossedRowsTPC(120);
  esdTrackCuts->SetMinRatioCrossedRowsOverFindableClustersTPC(0.8);
  esdTrackCuts->SetMaxChi2PerClusterITS(36);
  esdTrackCuts->SetMaxFractionSharedTPCClusters(0.4);
  esdTrackCuts->SetMaxChi2TPCConstrainedGlobal(36);
  
  esdTrackCuts->SetEtaRange(-0.9,0.9);
  esdTrackCuts->SetPtRange(0.15, 1e10);

  //___________________________________________________________________________
  //___________________________________________________________________________
  
  //11h_AOD095
  //https://twiki.cern.ch/twiki/bin/viewauth/ALICE/PWGPPAODTrackCuts
  //AOD095 	LHC11h 	pass2 	94M 	v5-02-Rev-04 	AddTaskESDFilter.C 	
  //http://alisoft.cern.ch/viewvc/tags/v5-02-Rev-04/ANALYSIS/macros/AddTaskESDFilter.C?view=markup&root=AliRoot
  
  Bool_t AddTrackCutsLHC11h(AliAnalysisTaskESDfilter* esdfilter)
  // standard cuts with tight DCA cut, using cluster cut instead of crossed rows (a la 2010 default)
  AliESDtrackCuts* esdTrackCutsH2Cluster = AliESDtrackCuts::GetStandardITSTPCTrackCuts2011(kTRUE, 0);
  // 1024, 1<< 10 // tight DCA cuts
  trackFilter->AddCuts(esdTrackCutsH2Cluster);
  //___________________________________________________________________________
  //___________________________________________________________________________
  */

  const TString runinfo = GetRunInfo();

  Int_t nshift = 10;
  if(runinfo.Contains("10h")){//2010
    nshift = 5; //10h_AOD049
  }
  else if(runinfo.Contains("11h")){//2011
    nshift = 10; //11h_AOD095
  }
  else{
    printf("AliAnalysisTaskAODpkRaa::CutTrack wrong run!! %s %d\n", runinfo.Data(), fEvent->GetRunNumber()); exit(1);
  }

  fFilterBit->Fill(nshift);

  const Int_t ibit = BIT(nshift);

  nsel = 0;
  for(Int_t ii=0; ii<fEvent->GetNumberOfTracks(); ii++){
    typetrack *trackptr = fEvent->GetTrack(ii);
    if(trackptr->TestFilterBit(ibit)){
      outtrk[nsel]=trackptr;
      nsel++;
    }
  }
}
